import re
import pickle
from argparse import ArgumentParser


NUMBER_PATTERN = re.compile(r'^(\d+,)*\d+(\.\d+){0,1}$')
PUNCT_PATTERN = re.compile(r'^\"|\'|-|--|,|\.|\(|\)|\?|\!|\[|\]|:|“|”|、|。|《|》$')

def is_number(word):
    return NUMBER_PATTERN.match(word) is not None

def is_punct(word):
    return PUNCT_PATTERN.match(word) is not None


def do_index(input_file, output_file):
    word2id = {}
    id2word = []

    input_file = open(input_file, 'r')
    output_file = open(output_file, 'w')

    for sent in input_file:
        words = sent.strip().lower().split()
        id_data = []
        for word in words:
            if (not is_number(word)) and (not is_punct(word)):
                if word not in word2id:
                    word2id[word] = len(word2id)
                    id2word.append(word)
                id_data.append(str(word2id[word]))
        output_file.write(' '.join(id_data) + '\n')

    input_file.close()
    output_file.close()

    return word2id, id2word


def main(option):
    word2id, id2word = do_index(option.source_file_in, option.source_file_out)
    src_vocab = { 'word2id': word2id, 'id2word': id2word }
    pickle.dump(src_vocab, open(option.source_vocab, 'wb'))

    word2id, id2word = do_index(option.target_file_in, option.target_file_out)
    tgt_vocab = { 'word2id': word2id, 'id2word': id2word }
    pickle.dump(tgt_vocab, open(option.target_vocab, 'wb'))


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--source_file_in', type=str, default='data/fbis.en.10k')
    parser.add_argument('--target_file_in', type=str, default='data/fbis.zh.10k')
    parser.add_argument('--source_file_out', type=str, default='output/fbis.en.10k')
    parser.add_argument('--target_file_out', type=str, default='output/fbis.zh.10k')
    parser.add_argument('--source_vocab', type=str, default='output/fbis.en.pkl')
    parser.add_argument('--target_vocab', type=str, default='output/fbis.zh.pkl')
    # parser.add_argument('--source_file_in', type=str, default='data/test.gr.txt')
    # parser.add_argument('--target_file_in', type=str, default='data/test.en.txt')
    # parser.add_argument('--source_file_out', type=str, default='output/test.gr.txt')
    # parser.add_argument('--target_file_out', type=str, default='output/test.en.txt')
    # parser.add_argument('--source_vocab', type=str, default='output/vocab_src.pkl')
    # parser.add_argument('--target_vocab', type=str, default='output/vocab_tgt.pkl')
    option = parser.parse_args()

    main(option)
